"""
将 MIDOG 2021 中  tif 转换成 yolo 系列需要的数据集

https://midog.deepmicroscopy.org/download-dataset/

"""
import json
import os
import numpy as np
from PIL import Image
import cv2
from tqdm import tqdm

if __name__ == '__main__':
    midog = 'MIDOG2021'
    wsi_path = f"/media/hsmy/wanghao_18T/dataset/{midog}/tiff_gt/"
    p = f"/media/hsmy/wanghao_18T/dataset/{midog}/classify_40_224/n/"
    n = f"/media/hsmy/wanghao_18T/dataset/{midog}/classify_40_224/p/"
    os.makedirs(p, exist_ok=True)
    os.makedirs(n, exist_ok=True)
    midog_json = f"/media/hsmy/wanghao_18T/dataset/{midog}/MIDOG.json"
    data = json.load(open(midog_json))

    patch_size = 224  # 训练 image 大小

    images_dict = {image["file_name"]: image["id"] for image in data["images"]}
    for wsi_file in tqdm(os.listdir(wsi_path)):
        image_id = images_dict[wsi_file]
        if image_id != 3:
            continue
        file_name = wsi_file.split('.')[0]
        file_path = os.path.join(wsi_path, wsi_file)
        img = cv2.imread(file_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        height, width, _ = img.shape

        annotations = [anno for anno in data["annotations"] if anno["image_id"] == image_id]

        if len(annotations) == 0:
            continue

        for index, anno in enumerate(annotations):
            bbox_arr = anno["bbox"]
            x_center = int(bbox_arr[0] + 25)
            y_center = int(bbox_arr[1] + 25)

            if x_center >= width or y_center >= height:
                continue

            patch_height, patch_width = patch_size, patch_size

            # 计算块的边界
            x_start = x_center - patch_width // 2
            x_end = x_start + patch_width
            y_start = y_center - patch_height // 2
            y_end = y_start + patch_height

            # 创建填充后的块
            patch = np.zeros((patch_height, patch_width, 3))

            # 计算图像与块之间的重叠区域
            x_start_clip = max(x_start, 0)
            x_end_clip = min(x_end, width)
            y_start_clip = max(y_start, 0)
            y_end_clip = min(y_end, height)

            # 计算填充块中的相对位置
            patch_x_start = max(0, -x_start)
            patch_x_end = patch_x_start + (x_end_clip - x_start_clip)
            patch_y_start = max(0, -y_start)
            patch_y_end = patch_y_start + (y_end_clip - y_start_clip)

            # 将重叠区域复制到填充块中
            patch[patch_y_start:patch_y_end, patch_x_start:patch_x_end, :] = img[y_start_clip:y_end_clip,
                                                                             x_start_clip:x_end_clip, :]
            img_item = Image.fromarray(np.uint8(patch))
            if anno["category_id"] == 1:
                img_item.save(f'{p}/{image_id}_{x_center}_{y_center}.png')
            else:
                img_item.save(f'{n}/{image_id}_{x_center}_{y_center}.png')
